calcBalanceTE <- function(data = NULL, actual.tr = NULL, vars.to.bal, vars.type, tr.type = "binary", tr.levels = 0:1, outcomes = NULL, id.vars = NULL, te.bl = "diff.means", te.complete = "diff.means", te.classical = "diff.means", seed = NULL, n.rands = 1, ...){

  if(tr.type == "binary"){
    if(!is.null(actual.tr)){
      bal.actual <- (balance.check(dat = data, vars.to.eval = vars.to.bal, vars.names = NULL, vars.type = vars.type, tr.var = actual.tr, 
                                   tr.names = NULL, tr.type = tr.type, graph.qq = FALSE, plot.file.qq = NULL, graph.pvals = FALSE,
                                   plot.file.pvals = NULL, text.xvalue = -.1, rnd = 3))$d2.p
    }else{
        bal.actual <- NULL
    }
  }
  if(tr.type == "continuous"){
    if(!is.null(actual.tr)){
      bal.actual <- (balance.check(dat = data, vars.to.eval = vars.to.bal, vars.names = NULL, vars.type = vars.type, tr.var = actual.tr, 
                                   tr.names = NULL, tr.type = tr.type, graph.qq = FALSE, plot.file.qq = NULL, graph.pvals = FALSE,
                                   plot.file.pvals = NULL, text.xvalue = -.1, rnd = 3))$pvals
    }else{
        bal.actual <- NULL
    }
  }
  
  store.te.bl <- store.bal.bl <- store.te.comp <- store.bal.comp <- store.te.clas <- store.bal.clas <- array(NA, n.rands)

  if(tr.type == "continuous"){
    store.bal.comp <- store.bal.clas <- store.bal.bl <- matrix(NA, n.rands, length(vars.to.bal))
    colnames(store.bal.comp) <- colnames(store.bal.clas) <- colnames(store.bal.bl) <- vars.to.bal
  }
  
  if(!is.null(seed)){
      set.seed(seed)
  }
  
  for(ii in 1:n.rands){
    if(ii%%5 == 0){
      print(paste("Randomization number ", ii, sep=""))
    }
    
    comp.assg <- sample(tr.levels, nrow(data), rep = TRUE)
    comp.data <- cbind(data, tr = comp.assg)    

    clas.idx <- list()
    clas.idx.tmp <- 1:nrow(data)
    clas.assg <- array(tr.levels[1], nrow(data))  ## minor to fix: don't just leave nonassigned as t=0
    for(jj in 1:length(tr.levels)){
      clas.idx[[jj]] <- sort(sample(clas.idx.tmp, floor(nrow(data)/length(tr.levels)), rep = FALSE))
      clas.assg[clas.idx[[jj]]] <- tr.levels[jj]
      clas.idx.tmp <- clas.idx.tmp[!(clas.idx.tmp %in% unlist(clas.idx))]
    }
    clas.data <- cbind(data, tr = clas.assg)

    out.b <- block(data, id.vars = id.vars, n.tr = length(tr.levels), ...)
    out.a <- assignment(out.b)

    ## create unique block IDs
    block.unique <- array(NA, nrow(data)) 
    for(id.idx in 1:nrow(data)){
      if(is.na(block.unique[id.idx])){
        for(ggg in 1:length(out.a$assg)){
          wh <- which(out.a$assg[[ggg]]== data[id.idx, id.vars])
          if(length(wh == 1)){
            block.row <- (wh %% nrow(out.a$assg[[ggg]]))
            if(block.row == 0){
              block.row <- nrow(out.a$assg[[ggg]])
            }
            block.unique[id.idx] <- paste("group", ggg, "block", block.row, sep="")
          }
        }
      }
    }  
      
    ## create vector of Tr values:
    bl.idx <- list()
    bl.assg <- array(tr.levels[1], nrow(data)) ## minor to fix: don't just leave nonassigned as t=0
    
    for(ttt in 1:length(tr.levels)){
      bl.tmp <- NULL
      for(ggg in 1:length(out.a$assg)){
        bl.tmp <- c(bl.tmp, as.vector(out.a$assg[[ggg]][, ttt]))
      }
      bl.idx[[ttt]] <- sort(bl.tmp[!is.na(bl.tmp)])
    }
    for(ttt in 1:length(tr.levels)){  ## bulky
      bl.assg[(which(data[, id.vars] %in% bl.idx[[ttt]]))] <- tr.levels[ttt]
    }
    
    ## combine  
    bl.id.vars <- dummify(as.factor(block.unique))
    bl.data <- cbind(data, tr = bl.assg, bl.id.vars)
      
    ## balance check
    bal.comp.obj <- balance.check(dat = comp.data, vars.to.eval = vars.to.bal, vars.names = NULL, vars.type = vars.type, tr.var = "tr", 
                                         tr.names = NULL, tr.type = tr.type, graph.qq = FALSE, plot.file.qq = NULL, graph.pvals = FALSE,
                                         plot.file.pvals = NULL, text.xvalue = -.1, rnd = 3)
    bal.clas.obj <- balance.check(dat = clas.data, vars.to.eval = vars.to.bal, vars.names = NULL, vars.type = vars.type, tr.var = "tr", 
                                         tr.names = NULL, tr.type = tr.type, graph.qq = FALSE, plot.file.qq = NULL, graph.pvals = FALSE, 
                                         plot.file.pvals = NULL, text.xvalue = -.1, rnd = 3)

    bal.bl.obj <- balance.check(dat = bl.data, vars.to.eval = vars.to.bal, vars.names = NULL, vars.type = vars.type, tr.var = "tr", 
                                       tr.names = NULL, tr.type = tr.type, graph.qq = FALSE, plot.file.qq = NULL, graph.pvals = FALSE,
                                       plot.file.pvals = NULL, text.xvalue = -.1, rnd = 3)
      
    if(tr.type == "binary"){

      ## store balance statistic
      store.bal.comp[ii] <- bal.comp.obj$d2.p
      store.bal.clas[ii] <- bal.clas.obj$d2.p
      store.bal.bl[ii] <- bal.bl.obj$d2.p
                                          
      if(te.complete == "diff.means"){
        store.te.comp[ii] <- mean(comp.data[comp.assg == 1, outcomes]) - mean(comp.data[comp.assg == 0, outcomes])      
      }
      if(te.classical == "diff.means"){
        store.te.clas[ii] <- mean(clas.data[clas.assg == 1, outcomes]) - mean(clas.data[clas.assg == 0, outcomes])      
      }
      if(te.bl == "diff.means"){
        store.te.bl[ii] <- mean(bl.data[bl.assg == 1, outcomes]) - mean(bl.data[bl.assg == 0, outcomes])      
      }
           
      if(te.complete == "regcoef"){
        if(nlevels(as.factor(data[, outcomes]))==2){
          store.te.comp[ii] <- as.numeric(glm(comp.data[, outcomes] ~ as.matrix(comp.data[, c("tr", vars.to.bal)]), family="binomial")$coefficients[2])
        }else{
        store.te.comp[ii] <- as.numeric(lm(comp.data[, outcomes] ~ as.matrix(comp.data[, c("tr", vars.to.bal)]))$coefficients[2])
        }
      }
      if(te.classical == "regcoef"){
        if(nlevels(as.factor(data[, outcomes]))==2){
          store.te.clas[ii] <- as.numeric(glm(clas.data[, outcomes] ~ as.matrix(clas.data[, c("tr", vars.to.bal)]), family="binomial")$coefficients[2])
        }else{
        store.te.clas[ii] <- as.numeric(lm(clas.data[, outcomes] ~ as.matrix(clas.data[, c("tr", vars.to.bal)]))$coefficients[2])
        }
      }
      if(te.bl == "regcoef"){
        if(nlevels(as.factor(data[, outcomes]))==2){
          store.te.bl[ii] <- as.numeric(glm(bl.data[, outcomes] ~ as.matrix(bl.data[, c("tr", vars.to.bal, colnames(bl.id.vars)[1:(length(colnames(bl.id.vars))-1)])]), family="binomial")$coefficients[2])
        }else{
        store.te.bl[ii] <- as.numeric(lm(bl.data[, outcomes] ~ as.matrix(bl.data[, c("tr", vars.to.bal)]))$coefficients[2])
        }
      }      
      
    }
      
    if(tr.type == "continuous"){
      
      ## store balance statistics
      store.bal.comp[ii, ] <- bal.comp.obj$pvals
      store.bal.clas[ii, ] <- bal.clas.obj$pvals
      store.bal.bl[ii, ] <- bal.bl.obj$pvals   
    }  
  }  

  return(list(bal.comp = store.bal.comp, bal.clas = store.bal.clas, bal.bl = store.bal.bl, bal.actual = bal.actual, te.comp = store.te.comp, te.clas = store.te.clas, te.bl = store.te.bl))

}
